
#pragma once
#include <eigen3/Eigen/Core>
#include <pcl/common/transforms.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/ModelCoefficients.h>

#include <gtsam/geometry/Rot3.h>
#include <gtsam/geometry/Pose3.h>
#include <gtsam/slam/PriorFactor.h>
#include <gtsam/slam/BetweenFactor.h>
#include <gtsam/nonlinear/NonlinearFactorGraph.h>
#include <gtsam/nonlinear/LevenbergMarquardtOptimizer.h>
#include <gtsam/nonlinear/GaussNewtonOptimizer.h>
#include <gtsam/nonlinear/Values.h>
#include <gtsam/inference/Symbol.h>
#include <gtsam/nonlinear/ISAM2.h>

#include "../3rdparty/sophus/se3.hpp"

#include "registration_interface.hpp"

namespace TESTICP
{
    class opt_ICP_GTSAM : public RegistrationInterface
    {
    public:
        EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
        opt_ICP_GTSAM(const YAML::Node& node);
        ~opt_ICP_GTSAM();
        void setInputTarget(const CloudPtr& target) override;
        void setInputSource(const CloudPtr& source) override;
        bool align( CloudType& output_ptr, const Eigen::Matrix4f& predict_pose = Eigen::Matrix4f::Identity()) override;

        float getFitnessScore();
        inline Eigen::Matrix4f getFinalTransformation() { return final_pose; }
    private:
        CloudPtr target_ptr, source_ptr;
        Eigen::Matrix4f final_pose;
        int max_iterations;
        float max_coresspoind_dis;
        bool use_point2plane;


        pcl::KdTreeFLANN<PointType>::Ptr kdtree_flann;
        
        // gtsam
        int laserCloudDSNum = 0;
        gtsam::ISAM2 *isam;
        gtsam::Values isamCurrentEstimate;
        gtsam::Point3 PointType2Point3(PointType point_in);
        gtsam::Pose3 trans2gtsamPose(Eigen::Matrix4f& transformIn);
        Eigen::Matrix4f pose2Matrix(const gtsam::Pose3& pose);
    };

    // 这个是点到平面的ICP方法，适用于对局部地图的配准。
    class Point2PlaneRigidTransformationFactor1: public gtsam::NoiseModelFactor1<gtsam::Pose3>
    {
    public:
        Point2PlaneRigidTransformationFactor1( gtsam::Key key,const gtsam::Point3 &P, const pcl::PointXYZI &Q,
                                            const gtsam::SharedNoiseModel &noise_model):
            gtsam::NoiseModelFactor1<gtsam::Pose3>(noise_model,key), P_(P), Q_(Q){}
        
        virtual ~Point2PlaneRigidTransformationFactor1(){}

        virtual gtsam::Vector evaluateError(const gtsam::Pose3& T,boost::optional<gtsam::Matrix&> H = boost::none)const override
        {

            gtsam::Matrix36 A_i;
            // 看他的原代码可以知道，前三维是旋转，后三维是平移(前三维我们是可以使用的)
            const gtsam::Point3 Q_i = T.transformFrom(P_, A_i); // 这个A是对p求雅克比矩阵，但是我们需要的是对位姿求雅克比矩阵
            gtsam::Vector3 n = (gtsam::Vector(3) << Q_.x, Q_.y, Q_.z).finished(); 
            gtsam::Vector3 temp_A_1 = n.transpose() * A_i.block(0, 0, 3, 3);
            gtsam::Matrix16 A =(gtsam::Matrix(1, 6) << temp_A_1[0], temp_A_1[1], temp_A_1[2], n[0], n[1], n[2]).finished();
            double error_i = Q_i.x() * Q_.x + Q_i.y() * Q_.y + Q_i.z() * Q_.z + Q_.intensity; // 这里计算修改为点到平面的误差，就可以
            if(H) *H = A; // 返回雅克比矩阵
            // std::cout << "The Error is : " << temp_A_1.transpose() << " : " << n.transpose()  << " : " << error_i << std::endl;
            return (gtsam::Vector(1) << error_i).finished();
        }

    private:
        gtsam::Point3 P_;
        pcl::PointXYZI Q_;
    };

    // 这个是点到点的ICP方法，适用于两个点云之间的配准
    class Point2PointRigidTransformationFactor: public gtsam::NoiseModelFactor1<gtsam::Pose3>
    {
    public:
        Point2PointRigidTransformationFactor( gtsam::Key key,const std::vector<gtsam::Point3> &P,
                                        const std::vector<gtsam::Point3>&Q,
                                        const gtsam::SharedNoiseModel &noise_model):
            gtsam::NoiseModelFactor1<gtsam::Pose3>(noise_model,key), P_(P), Q_(Q) {}
        
        virtual ~Point2PointRigidTransformationFactor() {}

        virtual gtsam::Vector evaluateError(const gtsam::Pose3& T,boost::optional<gtsam::Matrix&> H = boost::none)const override
        {
            gtsam::Matrix A = gtsam::Matrix::Zero(P_.size() * 3, 6);
            gtsam::Vector error = gtsam::Vector::Zero(P_.size() * 3);
            for(int i = 0; i < P_.size(); ++i)
            {
                gtsam::Matrix36 A_i;
                const gtsam::Point3 Q_i = T.transformFrom(P_[i], A_i);
                gtsam::Point3 error_i = Q_i - Q_[i];
                error.segment(3 * i, 3) = error_i;
                A.block(3 * i, 0, 3, 6) = A_i;
            }
            if(H)
            {
                *H = A;
            }
            return error;
        }

    private:
        std::vector<gtsam::Point3> P_;
        std::vector<gtsam::Point3> Q_;
    };


    opt_ICP_GTSAM::opt_ICP_GTSAM(const YAML::Node& node)
        : kdtree_flann(new pcl::KdTreeFLANN<PointType>) {
        
        max_iterations = node["max_iter"].as<int>();
        max_coresspoind_dis = node["max_corr_dist"].as<float>();
        use_point2plane = node["use_point2plane"].as<bool>();

        gtsam::ISAM2Params parameters;
        isam = new gtsam::ISAM2(parameters);
    }

    opt_ICP_GTSAM::~opt_ICP_GTSAM() { }

    void opt_ICP_GTSAM::setInputTarget(const CloudPtr& target) {
        target_ptr = target;
        kdtree_flann->setInputCloud(target);
        // std::cout << "the target_ptr size is : " << target_ptr->size() << std::endl;
    }

    void opt_ICP_GTSAM::setInputSource(const CloudPtr& source) {
        source_ptr = source;
        // std::cout << "the source_ptr size is : " << source_ptr->size() << std::endl;
    }

    bool opt_ICP_GTSAM::align( CloudType& output_ptr, const Eigen::Matrix4f& predict_pose) {
        CloudPtr transform_cloud(new CloudType());
        Eigen::Matrix4f T = predict_pose;
        gtsam::NonlinearFactorGraph graph;
        gtsam::Values initial;
        gtsam::SharedNoiseModel noise_model_new = gtsam::noiseModel::Isotropic::Sigma(1 , 0.1); // 考虑这个噪声的维度，这个是距离误差
        const gtsam::Key key = gtsam::symbol_shorthand::A(0);
        
        if(use_point2plane){
            // 构建点到平面的残差
            for (int i = 0; i < max_iterations; ++i) {
                // 清空数组
                std::cout << "the " << i  << " strat !!" << std::endl;
                
                transform_cloud->resize(0);
                pcl::transformPointCloud(*source_ptr, *transform_cloud, T);
                graph.resize(0);
                initial.clear();
                
                for (int j = 0; j < transform_cloud->size(); ++j) { // 遍历每个点，先构建点到平面的残差
                    PointType pointSel;
                    pcl::PointXYZI coeff;
                    PointType temp_point = transform_cloud->points[j];
                    if (!pcl::isFinite(temp_point))
                        continue;
                    
                    std::vector<float> res_dis;
                    std::vector<int> indices;
                    kdtree_flann->nearestKSearch(temp_point, 5, indices, res_dis);
                    if (res_dis[4] > max_coresspoind_dis)
                        continue;
                    Eigen::Matrix<float, 5, 3> matA0;
                    Eigen::Matrix<float, 5, 1> matB0;
                    Eigen::Vector3f matX0;

                    matA0.setZero();
                    matB0.fill(-1);
                    matX0.setZero();

                    for (int k = 0; k < 5; k++) {
                        matA0(k, 0) = target_ptr->points[indices[k]].x;
                        matA0(k, 1) = target_ptr->points[indices[k]].y;
                        matA0(k, 2) = target_ptr->points[indices[k]].z;
                    }
                    // 假设平面方程为ax+by+cz+1=0，这里就是求方程的系数abc，d=1
                    matX0 = matA0.colPivHouseholderQr().solve(matB0);

                    // 平面方程的系数，也是法向量的分量
                    float pa = matX0(0, 0);
                    float pb = matX0(1, 0);
                    float pc = matX0(2, 0);
                    float pd = 1;

                    // 单位法向量
                    float ps = sqrt(pa * pa + pb * pb + pc * pc);
                    pa /= ps; pb /= ps; pc /= ps; pd /= ps;

                    // 检查平面是否合格，如果5个点中有点到平面的距离超过0.2m，那么认为这些点太分散了，不构成平面
                    bool planeValid = true;
                    for (int k = 0; k < 5; k++) {
                        if (fabs(pa * target_ptr->points[indices[k]].x +
                                pb * target_ptr->points[indices[k]].y +
                                pc * target_ptr->points[indices[k]].z + pd) > 0.2) {
                            planeValid = false;
                            break;
                        }
                    }
                    if (!planeValid)
                        continue;
                    // 当前激光帧点到平面距离
                    float pd2 = pa * temp_point.x + pb * temp_point.y + pc * temp_point.z + pd;
                    float s = 1 - 0.9 * fabs(pd2) / sqrt(sqrt(temp_point.x * temp_point.x
                                + temp_point.y * temp_point.y + temp_point.z * temp_point.z));
                    // 点到平面垂线单位法向量（其实等价于平面法向量）
                    coeff.x = pa;
                    coeff.y = pb;
                    coeff.z = pc;
                    coeff.intensity = pd;
                    if (s > 0.1) {
                        gtsam::Point3 temp_P = PointType2Point3(temp_point);
                        graph.add(Point2PlaneRigidTransformationFactor1(key, temp_P, coeff, noise_model_new));
                    } else {
                        continue;
                    }
                }
                gtsam::Pose3 initial_pose = trans2gtsamPose(T);

                initial.insert(key, initial_pose);

                gtsam::LevenbergMarquardtOptimizer optimizer(graph, initial); // params_gn
                const gtsam::Values optimized = optimizer.optimize();
                gtsam::Pose3 transformed = optimized.at<gtsam::Pose3>(key); // 得到优化好的姿态
                transformed.rotation().print();
                transformed.translation().print();

                T = pose2Matrix(transformed) * T;
                
            }
        } else {
            // 构建点到点的残差
            std::vector<gtsam::Point3> P;
            std::vector<gtsam::Point3> Q;
            for (int i = 0; i < max_iterations; ++i) {
                std::cout << "the " << i  << " strat !!" << std::endl;
                transform_cloud->resize(0);
                pcl::transformPointCloud(*source_ptr, *transform_cloud, T);
                
                for (int j = 0; j < transform_cloud->size(); ++j) {
                    PointType temp_point = transform_cloud->points[j];
                    if (!pcl::isFinite(temp_point))
                        continue;
                    std::vector<float> res_dis;
                    std::vector<int> indices;
                    kdtree_flann->nearestKSearch(temp_point, 1, indices, res_dis);
                    if (res_dis.front() > max_coresspoind_dis)
                        continue;
                    P.push_back(PointType2Point3(temp_point));
                    Q.push_back(PointType2Point3(target_ptr->points[indices.front()]));

                }
                // std::cout << "the P.size() is " << P.size() << std::endl;
                gtsam::SharedNoiseModel noise_model = gtsam::noiseModel::Isotropic::Sigma(3 * P.size(), 1);
                graph.add(Point2PointRigidTransformationFactor(key, P, Q, noise_model));

                gtsam::Pose3 initial_pose = trans2gtsamPose(T);
                initial.insert(key, initial_pose);
                
                // 不同优化器（一共有三种）
                // gtsam::LevenbergMarquardtParams params_gn;
                // gtsam::LevenbergMarquardtOptimizer optimizer(graph, initial, params_gn);
                gtsam::GaussNewtonParams params_gn;
                gtsam::GaussNewtonOptimizer optimizer(graph, initial, params_gn);
                const gtsam::Values optimized = optimizer.optimize();
                gtsam::Pose3 transformed = optimized.at<gtsam::Pose3>(key); // 得到优化好的姿态
                transformed.rotation().print();
                transformed.translation().print();

                T = pose2Matrix(transformed) * T;

                P.clear();
                Q.clear();
                graph.resize(0);
                initial.clear();
            }
        }


        final_pose = T;
        pcl::transformPointCloud(*source_ptr, output_ptr, final_pose);

        return true;
    }

    gtsam::Point3 opt_ICP_GTSAM::PointType2Point3(PointType point_in) {
        gtsam::Point3 point_out(point_in.x, point_in.y, point_in.z);
        return point_out;
    }

    gtsam::Pose3 opt_ICP_GTSAM::trans2gtsamPose(Eigen::Matrix4f& transformIn) {
        Eigen::Matrix3f rotationMatrix = transformIn.block<3, 3>(0, 0);
        Eigen::Vector3f translation = transformIn.block<3, 1>(0, 3);

        // 将旋转矩阵转换为旋转角度
        Eigen::Vector3f euler = rotationMatrix.eulerAngles(0, 1, 2);

        return gtsam::Pose3(gtsam::Rot3::RzRyRx(euler(0), euler(1), euler(2)),
                gtsam::Point3(translation(0), translation(1), translation(2)));
    }

    Eigen::Matrix4f opt_ICP_GTSAM::pose2Matrix(const gtsam::Pose3& pose) {
        Eigen::Matrix4f transform = Eigen::Matrix4f::Identity();

        // 创建旋转矩阵
        Eigen::AngleAxisf rollAngle(pose.rotation().roll(), Eigen::Vector3f::UnitX());
        Eigen::AngleAxisf pitchAngle(pose.rotation().pitch(), Eigen::Vector3f::UnitY());
        Eigen::AngleAxisf yawAngle(pose.rotation().yaw(), Eigen::Vector3f::UnitZ());
        Eigen::Quaternionf rotationQuaternion = yawAngle * pitchAngle * rollAngle;
        Eigen::Matrix3f rotationMatrix = rotationQuaternion.matrix();

        // 设置平移向量
        transform(0, 3) = pose.translation().x();
        transform(1, 3) = pose.translation().y();
        transform(2, 3) = pose.translation().z();

        // 设置旋转矩阵
        transform.block<3, 3>(0, 0) = rotationMatrix;

        return transform;
    }

    float opt_ICP_GTSAM::getFitnessScore() {
        float max_range = std::numeric_limits<float>::max();
        float score = 0.f;

        CloudPtr transform_cloud(new CloudType());
        pcl::transformPointCloud(*source_ptr, *transform_cloud, final_pose);
        std::vector<int> nn_indices(1);
        std::vector<float> nn_dists(1);

        int nr = 0;

        for (size_t i = 0; i < transform_cloud->size(); ++i) {
            kdtree_flann->nearestKSearch(transform_cloud->points[i], 1, nn_indices, nn_dists);
            if (nn_dists.front() <= max_range) {
                score += nn_dists.front();
                nr++;
            }
        }
        if (nr > 0)
            return score / static_cast<float>(nr);
        else
            return (std::numeric_limits<float>::max());
    }
    
}


















